import torch
from torchvision import transforms
from PIL import Image
数据预处理 = transforms.Compose([
    transforms.Resize(256),  # 将图像大小调整为 256x256
    transforms.CenterCrop(224),  # 从中心裁剪出 224x224 的图像
    transforms.ToTensor(),  # 将图像转换为张量
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # 归一化
])
# 加载模型
模型 = torch.load('模型文件.pth',weights_only=False)
模型.eval()  # 将模型设置为评估模式
# 指定设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
模型.to(device)
# 读取图像
for i in range(1, 6):
    if i <= 2:
        image_path = str(i) + '.jpg'
    else:
        image_path = str(i) + '.png'
    #image_path = '2.jpg'  # 替换为你的图像路径
    image = Image.open(image_path).convert('RGB')  # 确保图像是RGB格式
    # 预处理图像
    image = 数据预处理(image).unsqueeze(0)  # 添加批量维度
    image = image.to(device)
    # 进行预测
    with torch.no_grad():
        output =模型(image)
        _, predicted = torch.max(output, 1)
    # 获取类别名称
    # CIFAR-10 的类别名称
    class_names = ['飞机', '汽车', '鸟', '猫', '鹿', '狗', '青蛙', '马', '船', '货车']
    predicted_class = class_names[predicted.item()]
    print(f"第{i}张图片识别结果是: {predicted_class}， 预测概率是: {output[0][predicted.item()]:.2f}")
    # print(f'图片识别结果是: {predicted_class}')